__all__ = ['parse_command_line', 'log_config', 'init', 'schedule_tasks', 'clean']


import argparse
import base64
import datetime
import logging
import os
import re
import sys
import time
import types

import util


__author__  = 'ZL <837806295@qq.com>'
__product__ = 'splitmhtml'
__version__ = '2.0.10'
__date__    = '2023/04/08'


log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())  # 如果宿主未配置日志处理器，则丢弃日志

args = types.SimpleNamespace()

truncate_length = 32
sequence_width = 5


def get_boundary_of_mhtml(data):
    """获取mhtml文件的分段分隔边界
    """

    if not isinstance(data, bytes) or not data:
        log.error('get boundary of mhtml ... failed | illegal input')
        return b''

    m = re.search(rb'Content\-Type\:.*?boundary\=\s*\"(.+?)\"\s+',
                  data,
                  re.DOTALL | re.IGNORECASE)
    if m and m.group() and m.group(1):
        boundary = m.group(1)
        log.debug('get boundary of mhtml ... successed')
    else:
        boundary = b''
        log.error('get boundary of mhtml ... failed')

    return boundary

def get_all_parts(data, boundary):
    """获取mhtml文件的所有分段
    """

    if (not isinstance(data, bytes)
            or not data
            or not isinstance(boundary, bytes)
            or not boundary):
        log.error('get all parts ... failed | illegal input')
        return []

    boundaries = []

    matchs = re.finditer(rb'\r\n\-\-%b\r\n' % re.escape(boundary),
                         data)
    for m in matchs:
        if m and m.group():
            boundaries.append((m.start(), m.end()))
        else:
            log.error('get all parts ... failed | boundary')
            return []

    m = re.search(rb'\r\n\-\-%b\-\-\r\n$' % re.escape(boundary),
                  data)
    if m and m.group():
        boundaries.append((m.start(), m.end()))
    else:
        log.error('get all parts ... failed | boundary')
        return []

    parts = []
    part_begin = 0
    for b in boundaries:
        parts.append((part_begin, b[0]))
        part_begin = b[1]

    log.debug('get all parts ... successed')
    return parts

def get_part_head(data):
    """获取mhtml文件的分段的头部
    """

    if not isinstance(data, bytes) or not data:
        log.error('get part ... failed | illegal input')
        return b''

    index = data.find(b'\r\n\r\n')
    if index != -1:
        head = data[0 : index + 2]
        log.debug('get part ... successed')
    else:
        head = b''
        log.error('get part ... failed')

    return head

def get_part_type(data):
    """获取mhtml文件的分段的类型
    """

    if not isinstance(data, bytes) or not data:
        log.error('get part type ... failed | illegal input')
        return ''

    m = re.search(rb'Content\-Type\:\s*([^\s]+?)\s+',
                  data,
                  re.IGNORECASE)
    if m and m.group() and m.group(1):
        type_ = m.group(1)
        log.debug('get part type ... successed')
    else:
        type_ = b''
        log.error('get part type ... failed')

    return type_.decode('utf-8')

def get_part_encoding(data):
    """获取mhtml文件的分段的编码
    """

    if not isinstance(data, bytes) or not data:
        log.error('get part encoding ... failed | illegal input')
        return ''

    m = re.search(rb'Content\-Transfer\-Encoding\:\s*([^\s]+?)\s+',
                  data,
                  re.IGNORECASE)
    if m and m.group() and m.group(1):
        encoding = m.group(1)
        log.debug('get part encoding ... successed')
    else:
        encoding = b''
        log.error('get part encoding ... failed | illegal input')

    return encoding.decode('utf-8')

def get_part_body(data):
    """获取mhtml文件的分段的正式内容
    """

    if not isinstance(data, bytes) or not data:
        log.error('get part body ... failed | illegal input')
        return b''

    length = len(data)
    index = data.find(b'\r\n\r\n')
    if index != -1 and index + 4 < length:
        body = data[index + 4 : length]
        log.debug('get part body ... successed')
    else:
        body = b''
        log.error('get part body ... failed')

    return body

def save_data(path, data):
    """保存字节数据到文件
    """

    if (not isinstance(path, str)
            or not path
            or not isinstance(data, bytes)
            or not data):
        log.error('save data ... failed | illegal input')
        return False

    try:
        with open(path, 'wb') as f:
            f.write(data)
    except Exception:
        log.exception('save data ... exception')
        return False

    log.debug('save data ... successed')
    return True

def save_base64(path, data):
    """保存base64编码的数据到文件
    """

    if (not isinstance(path, str)
            or not path
            or not (isinstance(data, bytes) or isinstance(data, str))
            or not data):
        log.error('save base64 ... failed | illegal input')
        return False

    try:
        data = base64.standard_b64decode(data)
    except Exception:
        log.exception('save base64 ... exception')
        return False

    return save_data(path, data)

def is_image_mime(mime_type):
    """判断是不是图像类型
    """

    if isinstance(mime_type, str) and mime_type.lower().startswith('image/'):
        log.debug('is image mime ... yes')
        return True
    else:
        log.debug('is image mime ... no | mime type: %s', mime_type)
        return False

def is_html_part(mime_type):
    """判断mhtml文件的分段是不是HTML分段
    """

    if isinstance(mime_type, str) and mime_type.lower().startswith('text/html'):
        log.debug('is html part ... yes')
        return True
    else:
        log.debug('is html part ... no | mime type: %s', mime_type)
        return False

def get_image_part_extension(mime_type):
    """获取mhtml文件的图像分段类型的扩展名部分
    """

    if not isinstance(mime_type, str) or not mime_type:
        log.critical('get image part extension ... failed ... exit | illegal '
                     'input')
        sys.exit(1)

    extension = mime_type.replace('image/', '')

    log.debug('get image part extension ... done | extension: %s', extension)
    return extension

def get_all_img_label(html):
    """从HTML中获取所有的<img>标签
    """

    if not isinstance(html, str) or not html:
        log.error('get all img label ... failed')
        return []

    m = re.findall(r'<img[^>]*>', html, re.IGNORECASE)
    if m:
        log.debug('get all img label ... successed')
        return m
    else:
        log.warning('get all img label ... failed')
        return []

def get_imge_data_from_img(img):
    """从<img>标签中获取所编码的图像数据
    """

    if not isinstance(img, str) or not img:
        log.error('get imge data from img ... failed')
        return []

    m = re.search(r'src="data:(image/jpg);base64,([^"]*)"',
                  img,
                  re.IGNORECASE)
    if m and m.group() and m.group(1) and m.group(2):
        log.debug('get imge data from img ... successed | jpg')
        return [m.group(1), 'base64', m.group(2)]

    m = re.search(r'src="data:(image/png);base64,([^"]*)"',
                  img,
                  re.IGNORECASE)
    if m and m.group() and m.group(1) and m.group(2):
        log.debug('get imge data from img ... successed | png')
        return [m.group(1), 'base64', m.group(2)]

    log.warning('get imge data from img ... failed')
    return []

def get_valid_parts(data):
    """获取去除了mhtml文件头部的其余部分
    """

    if not isinstance(data, bytes) or not data:
        log.error('get valid parts ... failed | illegal input')
        return []

    boundary = get_boundary_of_mhtml(data)
    if not boundary:
        log.error('get valid parts ... failed | boundary')
        return []

    parts = get_all_parts(data, boundary)
    count = len(parts)
    if count < 2:
        log.error('get valid parts ... failed | parts')
        return []

    log.debug('get valid parts ... successed')
    return parts[1:]

def parse_part(data):
    """分析mhtml分段返回分段信息
    """

    if not isinstance(data, bytes) or not data:
        log.error('parse part ... failed | illegal input')
        return []

    head = get_part_head(data)
    if not head:
        log.error('parse part ... failed | head')
        return []

    mime_type = get_part_type(head)
    if not mime_type:
        log.error('parse part ... failed | mime type')
        return []

    content_encoding = get_part_encoding(head)
    if not content_encoding:
        log.error('parse part ... failed | content encoding')
        return []

    content = get_part_body(data)
    if not content:
        log.error('parse part ... failed | content')
        return []

    log.debug('parse part ... successed')
    return [mime_type, content_encoding, content]

def save_part(out_dir,
              filename,
              mime_type,
              content_encoding,
              content):
    """保存mhtml文件的分段到文件
    """

    if (not isinstance(out_dir, str)
            or not out_dir
            or not os.path.isdir(out_dir)
            or not isinstance(filename, str)
            or not filename
            or not isinstance(mime_type, str)
            or not mime_type
            or not isinstance(content_encoding, str)
            or not content_encoding
            or not (isinstance(content, bytes) or isinstance(content, str))
            or not content):
        log.error('save part ... failed | illegal input')
        return False

    if is_image_mime(mime_type):
        extension = get_image_part_extension(mime_type)
        path = os.path.join(out_dir, '%s.%s' % (filename, extension))
    else:
        log.debug('save part ... ignored | not supported mime type: %s',
                  mime_type)
        return None

    if content_encoding == 'binary':
        return save_data(path, content)
    elif content_encoding == 'base64':
        return save_base64(path, content)
    else:
        log.warning('save part ... ignored | not supported encoding: %s',
                    content_encoding)
        return None

def split_single_mhtml(path, out_dir, name_func, name_info=None):
    """提取mhtml文件中的图片

    Args:
        path(str): 待提取的mhtml文件的完整路径
        out_dir(str): 存储提取文件的目录的完整路径
        name_func(func): 命名函数，提供分段的保存文件的名字

    Returns:
        文件为空时返回None
        出错返回False
        一切正常时返回True

    Return type: bool or None

    Raises:
        SystemExit
    """

    if (not isinstance(path, str)
            or not path
            or not os.path.isfile(path)
            or not isinstance(out_dir, str)
            or not out_dir
            or not os.path.isdir(out_dir)
            or not callable(name_func)):
        log.error('split single mhtml ... failed | illegal input')
        return False

    if os.path.getsize(path) <= 0:
        log.error('split single mhtml ... failed | empty file')
        return None

    try:
        with open(path, 'rb') as f:
            data = f.read()
    except Exception:
        log.exception('split single mhtml ... exception')
        return False

    parts = get_valid_parts(data)
    if not parts:
        log.error('split single mhtml ... failed | get parts')
        return False

    sequence = 0
    count = len(parts)
    for i in range(0, count):
        part = data[parts[i][0] : parts[i][1]]
        if not part:
            log.warning('split single mhtml ... ongoing | empty part')
            continue

        part2 = parse_part(part)
        if not part2:
            log.error('split single mhtml ... failed | parse part')
            return False

        if is_image_mime(part2[0]):
            filename = name_func(path, sequence, name_info)
            tristate = save_part(out_dir, filename, part2[0], part2[1],
                                 part2[2])
            if tristate is None:
                log.warning('split single mhtml ... ongoing | not supported '
                            'part')
                continue

            if not tristate:
                log.error('split single mhtml ... failed | save part')
                return False

            log.debug('split single mhtml ... ongoing | total: %d, done: %d, '
                      'image part sequence: %d',
                      count,
                      i + 1,
                      sequence)

            sequence += 1
        elif is_html_part(part2[0]):
            # 这儿假定HTML分段的内容是以UTF8格式编码
            html = part2[2].decode('utf-8')

            img_list = get_all_img_label(html)
            for img in img_list:
                image_data = get_imge_data_from_img(img)
                if not image_data:
                    log.warning('split single mhtml ... ongoing | not supported '
                                'img label')
                    continue

                filename = name_func(path, sequence, name_info)
                tristate = save_part(out_dir, filename, image_data[0],
                                     image_data[1],
                                     image_data[2])
                if tristate is None:
                    log.warning('split single mhtml ... ongoing | not supported '
                                'img label')
                    continue

                if not tristate:
                    log.error('split single mhtml ... failed | save part')
                    return False

                log.debug('split single mhtml ... ongoing | total: %d, done: '
                          '%d, img label sequence: %d',
                          count,
                          i + 1,
                          sequence)

                sequence += 1
        else:
            log.debug('split single mhtml ... ongoing | not target part')
            continue

    log.debug('split single mhtml ... successed')
    return True

def get_out_dir_for_file(path):
    """为处理单个mhtml文件生成输出目录

    如果命令行选项out_dir给了，就用给的，如果没有，则在待处理mhtml文件所在的目录创建一个子
    目录，用来存放提取的东西，子目录名由原mhtml的去扩展名的文件名的前32个字符与一个UUID用下
    划线连接而成
    """

    if (not isinstance(path, str)
            or not path
            or not os.path.isfile(path)):
        log.critical('get out dir for file ... failed ... exit | illegal input')
        sys.exit(1)

    if args.out_dir:
        out_dir = args.out_dir
    else:
        dir_, filename = os.path.split(os.path.splitext(path)[0])
        out_dir = os.path.join(dir_,
                               util.truncate(filename, truncate_length, True))

    try:
        os.makedirs(out_dir, exist_ok=True)
    except Exception:
        log.exception('get out dir for file ... exception ... exit')
        sys.exit(1)

    log.debug('get out dir for file ... successed')
    return out_dir

def name_for_file(path, sequence, name_info):
    """给分段要保存到的文件命名
    """

    if not isinstance(sequence, int) or sequence < 0:
        log.critical('name for file ... failed ... exit | illegal input')
        sys.exit(1)

    log.debug('name for file ... done | sequence: %d', sequence)
    return ('{:0>%d}' % sequence_width).format(sequence + 1)

def image_for_file():
    log.info('image for file ... ongoing')

    if util.has_extension2(args.path, args.dest_extension):
        tristate = split_single_mhtml(args.path,
                                      get_out_dir_for_file(args.path),
                                      name_for_file)
        if tristate is None:
            util.prompt('Ignore')
        else:
            util.prompt_assert(tristate, 'Split')
    else:
        util.prompt('Ignore')

    log.info('image for file ... done')

def split_single_mhtml_in_sandbox(path, out_dir, name_func, name_info=None):
    """提取单个mhtml文件

    提取单个mhtml文件，会拦截导致的任何异常，防止直接退出，以免影响别的任务
    """

    log.info('split single mhtml in sandbox ... ongoing | path: %s, '
             'out_dir: %s',
             path,
             out_dir)

    try:
        tristate = split_single_mhtml(path, out_dir, name_func, name_info)
    except Exception:
        log.exception('split single mhtml in sandbox ... exception | exception')
        return False
    except SystemExit as e:
        log.exception(
            'split single mhtml in sandbox ... exception | systemexit, '
            'exit code: %d',
            e.code)
        return False

    log.info('split single mhtml in sandbox ... done')
    return tristate

def get_out_dir_for_dir(path):
    """为处理目录下mhtml文件生成输出目录

    如果给了命令行选项out_dir，再看命令行选项share，如果share为真，则直接使用out_dir，如果
    share为假，则在out_dir给的目录下创建一个子目录，用来存放提取的东西，子目录名由原mhtml的
    去扩展名的文件名的前32个字符与一个UUID用下划线连接而成

    如果没给命令行选项out_dir，则在待处理mhtml文件所在的目录创建一个子目录，用来存放提取的东
    西，子目录名由原mhtml的去扩展名的文件名的前32个字符与一个UUID用下划线连接而成
    """

    if (not isinstance(path, str)
            or not path
            or not os.path.isfile(path)):
        log.critical('get out dir for dir ... failed ... exit | illegal input')
        sys.exit(1)

    dir_, filename = os.path.split(os.path.splitext(path)[0])

    if args.out_dir:
        if args.share:
            out_dir = args.out_dir
        else:
            out_dir = os.path.join(
                args.out_dir,
                util.truncate(filename, truncate_length, True))
    else:
        out_dir = os.path.join(dir_,
                               util.truncate(filename, truncate_length, True))

    try:
        os.makedirs(out_dir, exist_ok=True)
    except Exception:
        log.exception('get out dir for dir ... exception ... exit')
        sys.exit(1)

    log.debug('get out dir for dir ... successed')
    return out_dir

def name_for_dir(path, sequence, name_info):
    """给分段要保存到的文件命名

    当没有启用share命令行选项时，使用为单个mhtml提取内容时一样的命名方法

    当启用了share命令行选项时，按name命令行选项的取值来进行命名：

    original：使用“原文件名截断_UUID_内容序号”方式

    sequence：使用“日期_文件序号_内容序号”方式，其中文件序号与内容序号的最小占位字符宽度和
              起始序号可由命令行选项指定
    """

    if (not isinstance(path, str)
            or not path
            or not os.path.isfile(path)
            or not isinstance(sequence, int)
            or sequence < 0
            or not isinstance(name_info, dict)):
        log.critical('name for dir ... failed ... exit | illegal input')
        sys.exit(1)

    if ('sequence' not in name_info.keys()
            or not isinstance(name_info['sequence'], int)
            or name_info['sequence'] < 0):
        log.critical('name for dir ... failed ... exit | illegal input')
        sys.exit(1)

    file_sequence = name_info['sequence']
    dir_, filename = os.path.split(os.path.splitext(path)[0])

    if args.share:
        if args.name == 'original':
            if 'truncate' not in name_info.keys():
                name_info['truncate'] =  util.truncate(filename,
                                                       truncate_length,
                                                       True)

            name = ('%s_%s'
                    % (name_info['truncate'],
                       ('{:0>%d}' % sequence_width).format(sequence + 1)))
        elif args.name == 'sequence':
            now = datetime.datetime.now()
            name = (
                '%s_%s_%s'
                    % (
                        now.strftime('%Y%m%d'),
                        ('{:0>%d}' % args.file_width).format(
                            file_sequence + args.file_begin),
                        ('{:0>%d}' % args.content_width).format(
                            sequence + args.content_begin)
                    )
            )
        else:
            log.critical('name for dir ... failed ... exit | should not happen')
            sys.exit(1)
    else:
        name = ('{:0>%d}' % sequence_width).format(sequence + 1)

    log.debug('name for dir ... done | filename: %s, file sequence: %d, '
              'content sequence: %d',
              filename,
              file_sequence,
              sequence)
    return name

def image_for_dir():
    log.info('image for dir ... ongoing')

    total = [0]
    ignored = []
    failed = []

    try:
        sequence = 0
        for entry in os.listdir(args.path):
            entry_path = os.path.join(args.path, entry)

            if not os.path.isfile(entry_path):
                log.debug('image for dir ... ongoing | ignore: %s', entry_path)
                continue

            if not args.expose and util.is_hide(entry):
                log.info('image for dir ... ongoing | ignore: %s', entry_path)
                continue

            if not util.has_extension2(entry, args.dest_extension):
                log.debug('image for dir ... ongoing | ignore: %s', entry_path)
                continue

            tristate = split_single_mhtml_in_sandbox(
                entry_path,
                get_out_dir_for_dir(entry_path),
                name_for_dir,
                {'sequence': sequence})
            util.prompt_tristate(tristate, 'Split %s' % entry_path, False)
            util.make_statistics(entry_path, tristate, total, ignored, failed)

            if tristate:
                sequence += 1
    except Exception:
        log.exception('image for dir ... exception')
        util.prompt_assert(False, 'List %s' % args.path)

    util.prompt_batch(total[0], ignored, failed)

    log.info('image for dir ... done')

def image():
    log.info('image ... ongoing')

    if os.path.isfile(args.path):
        image_for_file()
    elif os.path.isdir(args.path):
        image_for_dir()

    log.info('image ... done')

def parse_command_line():
    parser = argparse.ArgumentParser(
        prog=__product__,
        description='拆解Chrome生成的MHTML文件提取内容')

    parser.add_argument('-v', '--version',
                        action='version',
                        version='%(prog)s ' + __version__)

    subparsers = parser.add_subparsers(
        title='subcommands',
        required=True,
        dest='subcommand',
        description='子命令集，由相应的子命令完成相应的功能')

    common = argparse.ArgumentParser(add_help=False)
    common.add_argument('-d', '--dest_extension',
                        type=util.arg_type_non_empty2,
                        nargs='+',
                        default=['.mht', '.mhtml'],
                        help='具有这些扩展名的文件才视为mhtml文件，不指定时默认会将具有'
                             '“.mht”和“.mhtml”两种扩展名的文件视为mhtml文件')
    common.add_argument('-o', '--out_dir',
                        type=util.arg_type_dir,
                        default=None,
                        help='默认会在被拆解的MHTML文件所在的目录下为每个mhtml文件创建'
                             '一个子目录来存放提取出来的文件，用此选项可以指定提取出来文'
                             '件的存放目录，当处理目录下的mhtml文件时，会在此目录下为每'
                             '个mhtml文件创建子目录，当处理单个mhtml文件时，直接在此目'
                             '录下存放提取出来的文件')
    common.add_argument('-s', '--share',
                        action='store_true',
                        help='当处理目录下的mhtml文件时，如果用outdir选项指定了提取出'
                             '来文件的存放目录，有此选项则不会单独为每个mhtml文件创建'
                             '子目录，而是所有的mhtml文件共用outdir选项指定的目录来'
                             '存放提取出来的文件')
    common.add_argument(
        '-n', '--name',
        choices=[
            'original',
            'sequence'
        ],
        default='sequence',
        help='当处理目录下的mhtml文件时，如果用outdir选项指定了提取出来文件的存放目录，用'
             'share选项指定了共享，此选项则用来指定为提取出来的文件命名的方式，original'
             '是“原文件名截断_UUID_内容序号”命名方式，sequence是“日期_文件序号_内容序号”'
             '命名方式，默认是sequence')
    common.add_argument('-b', '--file_begin',
                        type=util.arg_type_nonnegative,
                        default=1,
                        help='当name选项使用sequence值时，每个mhtml文件会对应一个序'
                             '号，此选项指定初始序号，默认为1')
    common.add_argument('-w', '--file_width',
                        type=util.arg_type_positive,
                        default=3,
                        help='当name选项使用sequence值时，每个mhtml文件会对应一个序'
                             '号，此选项指定此序号的占位宽度，序号大小不足位宽时前面会'
                             '补数字0，默认是3')
    common.add_argument('-B', '--content_begin',
                        type=util.arg_type_nonnegative,
                        default=1,
                        help='当name选项使用sequence值时，每个mhtml文件提取出来的每个'
                             '内容会对应一个序号，此选项指定初始序号，默认为1')
    common.add_argument('-W', '--content_width',
                        type=util.arg_type_positive,
                        default=3,
                        help='当name选项使用sequence值时，每个mhtml文件提取出来的每个'
                             '内容会对应一个序号，此选项指定此序号的占位宽度，序号大小不'
                             '足位宽时前面会补数字0，默认是3')
    common.add_argument('-e', '--expose',
                        action='store_true',
                        help='剥离目录下的MHTML文件时，默认会忽略目录下所有以“.”字符开'
                             '头的文件，有此选项则不忽略')

    # image子命令，用于提取图片
    image = subparsers.add_parser(
        'image',
        help='提取图片',
        parents=[util.arg_log(), util.arg_log_identifier(), common])
    image.add_argument('path',
                       type=util.arg_type_path,
                       help='待拆解的MHTML文件或者存有MHTML文件的目录的完整路径')

    global args
    args = parser.parse_args()

def log_config():
    if hasattr(args, 'log'):
        if hasattr(args, 'log_identifier') and args.log_identifier:
            return (
                '%s_v%s' % (__product__, __version__),
                args.log,
                args.log_identifier
            )
        else:
            return ('%s_v%s' % (__product__, __version__), args.log)

    return None

def init():
    log.info('program version: %s %s', __product__, __version__)
    log.debug('command line: %s', sys.argv)
    log.debug('parse args: %s', args)

def schedule_tasks():
    if args.subcommand == 'image':
        image()

def clean():
    pass
